import tensorflow as tf
from tensorflow.keras import layers, models
import pandas as pd
from sklearn.preprocessing import MultiLabelBinarizer

# 加载数据集
movies_df = pd.read_csv('movies.csv')

# 处理电影类型数据
genres = movies_df['genres'].apply(lambda x: x.split('|'))
mlb = MultiLabelBinarizer()
genres_encoded = mlb.fit_transform(genres)

# 将输入数据增加一个维度，以匹配卷积层期望的形状
genres_encoded = genres_encoded[:, :, None]

# 创建电影的特征输入
movie_input = tf.keras.Input(shape=(genres_encoded.shape[1], genres_encoded.shape[2]), name='Movie-Input')

# 创建网络层
# 一维卷积层
# Conv1D: 这是一维卷积层，用于在一维数据上执行卷积运算。filters=32意味着这层将学习32个不同的卷积核（或滤波器），每个卷积核都会检测特定的模式；kernel_size=3表示每个滤波器覆盖3个数据点，以便捕捉序列中的局部模式；activation='relu'是激活函数，用于引入非线性。
# MaxPooling1D: 池化层，用于降低数据的空间维度。pool_size=2意味着它将最大值取样，每两个数据点中取一个最大值。这有助于减少参数数量，缩小模型尺寸，防止过拟合，并在某种程度上保持特征的平移不变性。
# 第二个Conv1D和MaxPooling1D: 这是另一对卷积层和池化层，其作用类似于前一对层，但是这一次filters=64表示层将学习64个不同的特征表示。
# Flatten: 平展层，将上一层的输出转换为一维数组，以便可以将其传递给后续的全连接层。由于全连接层需要一维输入，所以这一步是必要的。
# Dense: 这是一个全连接层，units=64表示这一层由64个神经元组成。每个神经元将接收到前一层所有输出的加权和，经过非线性变换后输出。再次使用relu作为激活函数。
# 输出Dense层: 最后的全连接层，其神经元数量与电影类别的数量一致，即units=genres_encoded.shape[1]。它使用sigmoid激活函数，因为这个模型输出的是电影类别的概率，sigmoid函数的输出范围是(0, 1)，适用于二分类问题。

conv1 = layers.Conv1D(filters=32, kernel_size=3, activation='relu')(movie_input)
pool1 = layers.MaxPooling1D(pool_size=2)(conv1)
conv2 = layers.Conv1D(filters=64, kernel_size=3, activation='relu')(pool1)
pool2 = layers.MaxPooling1D(pool_size=2)(conv2)
flat = layers.Flatten()(pool2)
# 全连接层
fc1 = layers.Dense(units=64, activation='relu')(flat)
output = layers.Dense(units=genres_encoded.shape[1], activation='sigmoid')(fc1)

# 创建模型
model = models.Model(inputs=movie_input, outputs=output)
model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])

# 训练模型
tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir='./logs3', histogram_freq=1)
model.fit(genres_encoded, genres_encoded[:, :, 0], epochs=10, callbacks=[tensorboard_callback])

import numpy as np
# 推荐函数
def recommend(movie_id, top_n=5):
    # 我们需要使用np.expand_dims来增加输入数据的维度
    movie_vec = genres_encoded[movie_id]
    movie_vec = np.expand_dims(movie_vec, axis=0)  # 这将数据的形状从(18,)转变为(1, 18)
    prediction = model.predict(movie_vec)
    prediction = np.squeeze(prediction)  # 把结果的维度从(1, 18)降低回(18,)

    # 获取排序后的最高评分电影的索引
    sorted_indexes = prediction.argsort()[::-1][:top_n]

    # 返回对应的电影标题
    recommended_movie_ids = movies_df.iloc[sorted_indexes]['movieId']
    return movies_df[movies_df['movieId'].isin(recommended_movie_ids)]['title']


# import pdb
# pdb.set_trace()
# 推荐与movieId为1的电影相似的5部电影
print(recommend(1))
